import os
import SimpleITK as sitk
from PIL import Image
import numpy as np
import scipy.misc
import cv2
# from sklearn import
slice_number = 24
slice_size = 256

def resample_image_resize(patient_path, is_label=False, tar_size=(256, 256)):
    patient_name = patient_path.split('/')[-1]
    reader = sitk.ImageSeriesReader()
    series_IDs = reader.GetGDCMSeriesIDs(patient_path)
    file_reader = sitk.ImageFileReader()
    # dcm_series = reader.GetGDCMSeriesFileNames(patient_path)
    for series in series_IDs:
        series_file_names = reader.GetGDCMSeriesFileNames(patient_path, series)
        # 根据一个单张的dcm文件，读取这个series的metedata，即可以获取这个序列的描述符
        file_reader.SetFileName(series_file_names[0])
        file_reader.ReadImageInformation()
        series_description = file_reader.GetMetaData("0008|103e")
        if 't1' in series_description:
            itk_image = sitk.ReadImage(series_file_names)
            original_spacing = itk_image.GetSpacing()
            original_size = itk_image.GetSize()

            # 计算 reshape 后的 spacing
            reshape_spacing = [original_size[0] * original_spacing[0] / tar_size[0],
                               original_size[1] * original_spacing[1] / tar_size[1],
                               original_spacing[2]]

            # 计算 reshape 后的 size
            reshape_size = (tar_size[0], tar_size[1], original_size[2])

            # resample操作
            resample = sitk.ResampleImageFilter()
            resample.SetOutputSpacing(reshape_spacing)
            resample.SetSize(reshape_size)
            resample.SetOutputDirection(itk_image.GetDirection())
            resample.SetOutputOrigin(itk_image.GetOrigin())
            resample.SetTransform(sitk.Transform())
            resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())
            resize_img_tmp_t1 = resample.Execute(itk_image)
            # print(resize_img_tmp_t1.GetSpacing())
            resize_img_tmp_t1 = resample.Execute(resize_img_tmp_t1)
            data = sitk.GetArrayFromImage(resize_img_tmp_t1)
            print(data.shape)
            for i in range(data.shape[0]):
                slice = Image.fromarray(data[i, :, :])
                # cv2.imwrite('./BAI_RUI_YING/resize/t1_' + str(i) + '.png', np.uint8(slice))
                # scipy.misc.imsave('./BAI_RUI_YING/resize/t1_' + str(i) + '.png', slice)
                # slice.save('./BAI_RUI_YING/resize/t1_' + str(i) + '.png')

                slice *= 255  # 变换为0-255的灰度值
                im = Image.fromarray(slice)
                im = im.convert('L')  # 这样才能转为灰度图，如果是彩色图则改L为‘RGB’
                im.save('./BAI_RUI_YING/resize/t1_' + str(i) + '.png')
        if 't2' in series_description:
            itk_image = sitk.ReadImage(series_file_names)
            original_spacing = itk_image.GetSpacing()
            original_size = itk_image.GetSize()

            # 计算 reshape 后的 spacing
            reshape_spacing = [original_size[0] * original_spacing[0] / tar_size[0],
                               original_size[1] * original_spacing[1] / tar_size[1],
                               original_spacing[2]]

            # 计算 reshape 后的 size
            reshape_size = (tar_size[0], tar_size[1], original_size[2])

            # resample操作
            resample = sitk.ResampleImageFilter()
            resample.SetOutputSpacing(reshape_spacing)
            resample.SetSize(reshape_size)
            resample.SetOutputDirection(itk_image.GetDirection())
            resample.SetOutputOrigin(itk_image.GetOrigin())
            resample.SetTransform(sitk.Transform())
            resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())
            resize_img_tmp_t2 = resample.Execute(itk_image)
            data = sitk.GetArrayFromImage(resize_img_tmp_t2)
            print(data.shape)
            for i in range(data.shape[0]):
                slice = Image.fromarray(data[i,:,:])
                # cv2.imwrite('./BAI_RUI_YING/resize/t2_' + str(i) + '.png', np.uint8(slice))
                # slice.save(scipy.misc.('./BAI_RUI_YING/resize/t2_' + str(i) + '.png', slice))
                slice *= 255  # 变换为0-255的灰度值
                im = Image.fromarray(slice)
                im = im.convert('L')  # 这样才能转为灰度图，如果是彩色图则改L为‘RGB’
                im.save('./BAI_RUI_YING/resize/t2_' + str(i) + '.png')




    # reader.SetFileNames(dcm_series)
    # itk_image = reader.Execute()
    # original_spacing = itk_image.GetSpacing()
    # original_size = itk_image.GetSize()
    #
    # # 计算 reshape 后的 spacing
    # reshape_spacing = [original_size[0] * original_spacing[0] / tar_size[0],
    #                    original_size[1] * original_spacing[1] / tar_size[1],
    #                    original_spacing[2]]
    #
    # # 计算 reshape 后的 size
    # reshape_size = (tar_size[0], tar_size[1], original_size[2])
    #
    # # resample操作
    # resample = sitk.ResampleImageFilter()
    # resample.SetOutputSpacing(reshape_spacing)
    # resample.SetSize(reshape_size)
    # resample.SetOutputDirection(itk_image.GetDirection())
    # resample.SetOutputOrigin(itk_image.GetOrigin())
    # resample.SetTransform(sitk.Transform())
    # resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())
    #
    # if is_label:
    #     resample.SetInterpolator(sitk.sitkNearestNeighbor)
    # else:
    #     resample.SetInterpolator(sitk.sitkBSpline)
    #
    # resize_img_tmp = resample.Execute(itk_image)
    #
    # new = sitk.ReadImage()

    # 保存为： 病人名.nii.gz 格式
    # sitk.WriteImage(resize_img_tmp, os.path.join(patient_path, patient_name + '.nii.gz'))


def img_center_crop(image, crop_size):
    assert len(image.shape) == 3, 'invalid image size in sliding window'

    z_start, x_start, y_start = 0, 0, 0
    img_z, img_x, img_y = image.shape[0], image.shape[1], image.shape[2]
    crop_z, crop_x, crop_y = crop_size[0], crop_size[1], crop_size[2]
    # if img_x>set_size and img_y>set_size
    # x or y 一个比crop大 或者两个都大
    if img_x > crop_x and img_y > crop_y:
        starting = [int((crop_z - img_z) / 2), int((img_x - crop_x) / 2), int((img_y - crop_y) / 2)]
        z_start, x_start, y_start = starting[0], starting[1], starting[2]

    elif img_x > crop_x and img_y <= crop_y:
        starting = [int((crop_z - img_z) / 2), int((img_x - crop_x) / 2), int((crop_y - img_y) / 2)]
        z_start, x_start, y_start = starting[0], starting[1], 0
        y_up = int((crop_y - img_y) / 2)
        y_down = crop_y - img_y - y_up
        image = np.pad(image, ((0, 0), (0, 0), (y_up, y_down)), mode='constant')

    elif img_x <= crop_x and img_y > crop_y:
        starting = [int((crop_z - img_z) / 2), int((crop_x - img_x) / 2), int((img_y - crop_y) / 2)]
        z_start, x_start, y_start = starting[0], 0, starting[2]
        x_up = int((crop_x - img_x) / 2)
        x_down = crop_x - img_x - x_up
        image = np.pad(image, ((0, 0), (x_up, x_down), (0, 0)), mode='constant')

    img_crop = image[z_start: z_start + crop_size[0], x_start:x_start + crop_size[1],
               y_start: y_start + crop_size[2]]

    return img_crop

import nrrd 

def resample_image(itk_image, out_spacing=(1., 1., 1.), is_label=False):
    original_spacing = itk_image.GetSpacing()
    original_size = itk_image.GetSize()
    out_size = [int(np.round(original_size[0] * (original_spacing[0] / out_spacing[0]))),
                int(np.round(original_size[1] * (original_spacing[1] / out_spacing[1]))),
                int(np.round(original_size[2] * (original_spacing[2] / out_spacing[2])))]

    resample = sitk.ResampleImageFilter()
    resample.SetOutputSpacing(out_spacing)
    resample.SetSize(out_size)
    resample.SetOutputDirection(itk_image.GetDirection())
    resample.SetOutputOrigin(itk_image.GetOrigin())
    resample.SetTransform(sitk.Transform())
    resample.SetDefaultPixelValue(itk_image.GetPixelIDValue())

    if is_label:
        resample.SetInterpolator(sitk.sitkNearestNeighbor)
    else:
        resample.SetInterpolator(sitk.sitkBSpline)

    return resample.Execute(itk_image)

import zipfile
import os 

def unzip_downloaded_dicom(p):
    f = zipfile.ZipFile(p,'r') # 压缩文件位置
    all_files = f.namelist()

    for file in f.namelist():
        f.extract(file,"./")
    f.close()

    return all_files

def padding_image_array_size(image_array, out_size):
    img_z, img_x, img_y = image_array.shape[0], image_array.shape[1], image_array.shape[2]
    out_z, out_x, out_y = out_size[0], out_size[1], out_size[2]
    if out_z > img_z:
        z_up = int((out_z - img_z) / 2)
        z_down = out_z - img_z - z_up

        if out_x >= img_x and out_y >= img_y:  # 三个维度都是padding
            x_up = int((out_x - img_x) / 2)
            x_down = out_x - img_x - x_up
            y_up = int((out_y - img_y) / 2)
            y_down = out_y - img_y - y_up
            new_volume = np.pad(image_array, ((z_up, z_down), (x_up, x_down), (y_up, y_down)), mode='constant')
        else:
            new_volume = np.pad(image_array, (z_up, z_down), mode='constant')
            new_volume = img_center_crop(new_volume, (slice_number, slice_size, slice_size))
    else:
        # 把z轴crop为slice number
        z_start = int((img_z - out_z) / 2)
        image_array = image_array[z_start: z_start + out_size[0], :, :]
        if out_x >= img_x and out_y >= img_y:  # 俩个维度都是padding
            x_up = int((out_x - img_x) / 2)
            x_down = out_x - img_x - x_up
            y_up = int((out_y - img_y) / 2)
            y_down = out_y - img_y - y_up
            new_volume = np.pad(image_array,((0,0),(x_up, x_down), (y_up, y_down)), mode='constant')
        else:
            new_volume = img_center_crop(image_array, (24, 256, 256))

    return new_volume

if __name__ == '__main__':
    path = './huanhu_raw_data'
    new_path= './huanhu_raw_data_resampled'
    # tar = (256, 256)  # 只定义 x,y 方向的分辨率，由于数据是3D volume，所以z(slice)可以保持不变。
    #
    # resample_image(patient_path, is_label=False, tar_size=tar)
    # patient_name = patient_path.split('/')[-1]
    for label_path_name in os.listdir(path):
        if label_path_name == 'label':
            label_path = path+'/'+label_path_name
            for grade_path_name in os.listdir(label_path):
                grade_path = label_path+'/'+grade_path_name
                for patient_path_name in os.listdir(grade_path):
                    patient_path = grade_path+'/'+patient_path_name
                    new_patient_path = new_path+'/'+label_path_name+'/'+grade_path_name+'/'+patient_path_name
                    print(new_patient_path)
                    if not os.path.exists(new_patient_path):
                        os.makedirs(new_patient_path)
                                    
                    reader = sitk.ImageSeriesReader()
                    series_IDs = reader.GetGDCMSeriesIDs(patient_path)
                    file_reader = sitk.ImageFileReader()
                    # dcm_series = reader.GetGDCMSeriesFileNames(patient_path)
                    for series in series_IDs:
                        series_file_names = reader.GetGDCMSeriesFileNames(patient_path, series)
                        # 根据一个单张的dcm文件，读取这个series的metedata，即可以获取这个序列的描述符
                        file_reader.SetFileName(series_file_names[0])
                        file_reader.ReadImageInformation()
                        series_description = file_reader.GetMetaData("0008|103e")
                        if 't1' in series_description and 'tra' in series_description:
                            itk_image = sitk.ReadImage(series_file_names)
                            t1_raw_spacing = itk_image.GetSpacing()
                            resampled_image = resample_image(itk_image, out_spacing=(1., 1., 1.),is_label=False)  # itk_image.GetSize (x,y,z)
                            resampled_image = sitk.GetArrayFromImage(resampled_image)  # GetArrayFromImage (z,x,y)
                            image_resample = padding_image_array_size(resampled_image, out_size=(24, 256, 256))
                            itk_image_resample = sitk.GetImageFromArray(image_resample)
                            sitk.WriteImage(itk_image_resample, new_patient_path+'/t1.nii.gz')

                        if 't2' in series_description and 'tra' in series_description and 'dark-fluid' in series_description:
                            itk_image = sitk.ReadImage(series_file_names)
                            resampled_image = resample_image(itk_image, out_spacing=(1., 1., 1.),is_label=False)  # itk_image.GetSize (x,y,z)
                            resampled_image = sitk.GetArrayFromImage(resampled_image)  # GetArrayFromImage (z,x,y)
                            image_resample = padding_image_array_size(resampled_image, out_size=(24, 256, 256))
                            itk_image_resample = sitk.GetImageFromArray(image_resample)
                            sitk.WriteImage(itk_image_resample, new_patient_path + '/t2.nii.gz')


                    ## 读取分割数据
                    all_paths = os.listdir(patient_path)
                    for p in all_paths:
                        label_path_nrrd = os.path.join(patient_path, p)

                        if "PN0" not in label_path_nrrd:
                            continue

                        print(f"label_path_nrrd is {label_path_nrrd}")
                        label_dirs = os.listdir(label_path_nrrd)

                        for label_dir in label_dirs:
                            label_dir_nrrd = os.path.join(label_path_nrrd, label_dir)

                            for label in os.listdir(label_dir_nrrd):
                                if "T1WI-CE_t_1.nrrd.zip" in label:
                                    all_files = unzip_downloaded_dicom(os.path.join(label_dir_nrrd, label))
                                    print(all_files)
                                    data_path = all_files[0]

                                    data, options=nrrd.read(data_path)
                                    import matplotlib.pyplot as plt 

                                    data = data.transpose([2, 1, 0])
                                    print(data.shape)
                                    label_itk_image = sitk.GetImageFromArray(data)
                                    label_itk_image.SetSpacing(t1_raw_spacing)

                                    resampled_label_image = resample_image(label_itk_image, out_spacing=(1., 1., 1.),is_label=True)  # itk_image.GetSize (x,y,z)
                                    resampled_label_image = sitk.GetArrayFromImage(resampled_label_image)  # GetArrayFromImage (z,x,y)
                                    resampled_label_image = padding_image_array_size(resampled_label_image, out_size=(24, 256, 256))
                                    resampled_label_image = sitk.GetImageFromArray(resampled_label_image)
                                    sitk.WriteImage(resampled_label_image, new_patient_path+'/tumor.nii.gz')
                                    print(new_patient_path + '/tumor.nii.gz')


